package com.github.jiangwangyang.a2a.client;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import io.a2a.http.A2AHttpClient;
import io.a2a.http.A2AHttpResponse;
import io.a2a.spec.JSONRPCError;
import io.a2a.spec.StreamingEventKind;
import io.a2a.util.Utils;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

public class WebClientA2AHttpClient implements A2AHttpClient {

    public static final ParameterizedTypeReference<ServerSentEvent<String>> SSE_STRING_TYPE_REFERENCE = new ParameterizedTypeReference<>() {
    };

    private final WebClient webClient;

    public WebClientA2AHttpClient(WebClient webClient) {
        this.webClient = webClient;
    }

    public Flux<StreamingEventKind> requestFlux(@NonNull HttpMethod method,
                                                @NonNull String url,
                                                @NonNull Map<String, String> headers,
                                                @NonNull String body) {
        return webClient.method(method)
                .uri(url)
                .headers(httpHeaders -> headers.forEach(httpHeaders::set))
                .header("Content-Type", "application/json")
                .header("Accept", "text/event-stream")
                .bodyValue(body)
                .retrieve()
                .bodyToFlux(SSE_STRING_TYPE_REFERENCE)
                .map(ServerSentEvent::data)
                .map(message -> {
                    try {
                        JsonNode jsonNode = Utils.OBJECT_MAPPER.readTree(message);
                        if (jsonNode.has("error")) {
                            throw Utils.OBJECT_MAPPER.treeToValue(jsonNode.get("error"), JSONRPCError.class);
                        } else if (jsonNode.has("result")) {
                            JsonNode result = jsonNode.path("result");
                            return Utils.OBJECT_MAPPER.treeToValue(result, StreamingEventKind.class);
                        } else {
                            throw new IllegalArgumentException("Unknown message type");
                        }
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException(e);
                    }
                });
    }

    @Override
    public GetBuilder createGet() {
        return new GetBuilder() {
            private final Map<String, String> headers = new HashMap<>();
            private String url = null;

            @Override
            public A2AHttpResponse get() {
                return requestMono(HttpMethod.GET, url, headers, null);
            }

            @Override
            public CompletableFuture<Void> getAsyncSSE(Consumer<String> messageConsumer, Consumer<Throwable> errorConsumer, Runnable completeRunnable) {
                headers.put("Accept", "text/event-stream");
                return requestFuture(HttpMethod.GET, url, headers, null, messageConsumer, errorConsumer, completeRunnable);
            }

            @Override
            public GetBuilder url(String s) {
                url = s;
                return this;
            }

            @Override
            public GetBuilder addHeader(String name, String value) {
                headers.put(name, value);
                return this;
            }
        };
    }

    @Override
    public PostBuilder createPost() {
        return new PostBuilder() {
            private final Map<String, String> headers = new HashMap<>();
            private String url = null;
            private String body = null;

            @Override
            public A2AHttpResponse post() {
                return requestMono(HttpMethod.POST, url, headers, body);
            }

            @Override
            public CompletableFuture<Void> postAsyncSSE(Consumer<String> messageConsumer, Consumer<Throwable> errorConsumer, Runnable completeRunnable) {
                headers.put("Accept", "text/event-stream");
                return requestFuture(HttpMethod.POST, url, headers, body, messageConsumer, errorConsumer, completeRunnable);
            }

            @Override
            public PostBuilder body(String body) {
                this.body = body;
                return this;
            }

            @Override
            public PostBuilder url(String s) {
                url = s;
                return this;
            }

            @Override
            public PostBuilder addHeader(String name, String value) {
                headers.put(name, value);
                return this;
            }
        };
    }

    protected CompletableFuture<Void> requestFuture(
            @NonNull HttpMethod method, @NonNull String url, @NonNull Map<String, String> headers, @Nullable String body,
            @NonNull Consumer<String> messageConsumer, @NonNull Consumer<Throwable> errorConsumer, @NonNull Runnable completeRunnable) {
        return webClient.method(method)
                .uri(url)
                .headers(h -> headers.forEach(h::add))
                .bodyValue(Optional.ofNullable(body).orElse(""))
                .retrieve()
                .bodyToFlux(SSE_STRING_TYPE_REFERENCE)
                .map(ServerSentEvent::data)
                .map(msg -> {
                    messageConsumer.accept(msg);
                    return null;
                })
                .onErrorResume(throwable -> {
                    errorConsumer.accept(throwable);
                    return Mono.empty();
                })
                .doOnComplete(completeRunnable)
                .then()
                .toFuture();
    }

    protected A2AHttpResponse requestMono(@NonNull HttpMethod method,
                                          @NonNull String url,
                                          @NonNull Map<String, String> headers,
                                          @Nullable String body) {
        return webClient.method(method)
                .uri(url)
                .headers(h -> headers.forEach(h::add))
                .bodyValue(Optional.ofNullable(body).orElse(""))
                .retrieve()
                .toEntity(String.class)
                .map(StringA2AHttpResponse::new)
                .block();
    }

    record StringA2AHttpResponse(ResponseEntity<String> responseEntity) implements A2AHttpResponse {
        @Override
        public int status() {
            return responseEntity.getStatusCode().value();
        }

        @Override
        public boolean success() {
            return responseEntity.getStatusCode().is2xxSuccessful();
        }

        @Override
        public String body() {
            return responseEntity.getBody();
        }
    }

}
